import numpy as np
import matplotlib.pyplot as plt
from pettingzoo.mpe import simple_spread_v3
from matplotlib.patches import Circle

env = simple_spread_v3.parallel_env(N=3, local_ratio=0.5, max_cycles=25)
env.reset(seed=42)

def get_positions(env):
    obs = env.state()
    agent_num = len(env.agents)
    landmark_num = getattr(env, 'num_landmarks', None)
    if landmark_num is None and hasattr(env, 'env'):
        landmark_num = getattr(env.env, 'num_landmarks', 0)

    agent_pos = [obs[i * 4 : i * 4 + 2] for i in range(agent_num)]
    landmark_pos = [obs[agent_num * 4 + i * 2 : agent_num * 4 + i * 2 + 2] for i in range(landmark_num)]
    return np.array(agent_pos), np.array(landmark_pos)

agent_pos, landmark_pos = get_positions(env)

fig, ax = plt.subplots(figsize=(6, 6))
ax.set_xlim(-1.5, 1.5)
ax.set_ylim(-1.5, 1.5)
ax.set_title("Simple Spread v3 - Initial Positions")
ax.set_aspect('equal')

for i, pos in enumerate(landmark_pos):
    ax.add_patch(Circle(pos, 0.05, color='gray'))
    ax.text(pos[0], pos[1] + 0.07, f"Landmark {i}", ha='center', fontsize=8)

for i, pos in enumerate(agent_pos):
    ax.add_patch(Circle(pos, 0.05, color='blue'))
    ax.text(pos[0], pos[1] - 0.07, f"Agent {i}", ha='center', color='blue', fontsize=8)

plt.grid(True)
plt.show()
